import numpy as np
import cv2
from scipy.fftpack import dct, idct


def embed_watermark_dct(host_img, watermark, alpha=0.1):
    """
    宿主图像: 灰度图 (0-255)
    水印: 二值矩阵 (0/1)
    alpha: 嵌入强度
    """
    h, w = host_img.shape
    watermarked = host_img.copy().astype(np.float32)
    wm_size = watermark.size
    wm_flat = watermark.flatten()

    # 分块处理
    for i in range(0, h, 8):
        for j in range(0, w, 8):
            block = host_img[i:i + 8, j:j + 8]
            if block.shape != (8, 8): continue

            # DCT变换
            dct_block = dct(dct(block.T, norm='ortho').T, norm='ortho')

            # 在中频系数嵌入水印（示例：对角线位置）
            if wm_size > 0:
                idx = min(i // 8 * (w // 8) + j // 8, wm_size - 1)
                dct_block[4, 4] += alpha * (255 if wm_flat[idx] > 0 else -255)

            # IDCT逆变换
            watermarked[i:i + 8, j:j + 8] = idct(idct(dct_block.T, norm='ortho').T, norm='ortho')

    return np.clip(watermarked, 0, 255).astype(np.uint8)


def extract_watermark_dct(watermarked_img, wm_shape, alpha=0.1):
    h, w = watermarked_img.shape
    watermark = np.zeros(wm_shape).flatten()

    for i in range(0, h, 8):
        for j in range(0, w, 8):
            block = watermarked_img[i:i + 8, j:j + 8]
            if block.shape != (8, 8): continue

            dct_block = dct(dct(block.T, norm='ortho').T, norm='ortho')
            idx = i // 8 * (w // 8) + j // 8
            if idx < len(watermark):
                watermark[idx] = 1 if dct_block[4, 4] > 0 else 0

    return watermark.reshape(wm_shape)


# 生成测试数据
host = cv2.imread("srcGray.jpg", cv2.IMREAD_GRAYSCALE)

def GetWaterMark():
    # watermark = np.random.randint(0, 2, (32, 32))  # 32x32二值水印
    watermark = np.random.randint(0, 2, (32, 32))  # 32x32二值水印
    return watermark


watermark = GetWaterMark()

# 嵌入与提取
watermarked_img = embed_watermark_dct(host, watermark)
cv2.imwrite("embededWatermark.jpg", watermarked_img)
extracted_wm = extract_watermark_dct(watermarked_img, watermark.shape)

# 评估相似度
similarity = np.mean(watermark == extracted_wm)
print(f"水印相似度: {similarity * 100:.2f}%")